package correlation;
import java.util.Date;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;
import java.util.ListIterator;
import java.util.Set;

import kb.KnowledgeBase;
import dom.Alert;
import dom.HyperAlert;
import dom.Implication;
import dom.Predicate;

public class CorrelationEngine {
	public static class Key{
		String key;		
		Key(HyperAlert a, Predicate p){
			List<String> facts = p.getFacts();
			String alertInfo = "";
			for(String s : facts){
				alertInfo = a.getFact(s);
			}
			key = p.getName()+alertInfo;
		}
		@Override
		public String toString() {
			return "Key [key=" + key + "]";
		}
		
		public int hashCode(){
			return key.hashCode();
		}
		public boolean equals(Object o){
			if(o instanceof Key){
				return ((Key) o).key.equals(this.key);
			}
			return false;
		}
	}
	
	private static boolean debug = false;
	private static List<HyperAlertCorrelationNode> acn = new LinkedList<HyperAlertCorrelationNode>();
	private static HashMap<Key,List<HyperAlertCorrelationNode>> PredicateFactToNode = new HashMap<Key,List<HyperAlertCorrelationNode>>();//Finding Consequences
	private List<HyperAlert> hyperAlerts = new LinkedList<HyperAlert>();
	private List<Implication> implications = new LinkedList<Implication>();
	int gracePeriodHrs = 2; 
	
	public CorrelationEngine() {}
	
	public CorrelationEngine(List<Alert> alerts) { 
		initialize(alerts);
	}
	
	public void initialize(List<Alert> alerts){
		if(debug){
			System.out.println("Entered initialize() in Correlation Engine");
		}
		KnowledgeBase kb = KnowledgeBase.getInstance();
		kb.load();
		for(Alert a : alerts) {
			hyperAlerts.add( new HyperAlert(a.getSig_sid()
										  , a.getTimestamp()
										  , a.getIp_src(), a.getPort_src()
										  , a.getIp_dst(), a.getPort_dst()));
		}
		implications = kb.getImplications();
	}
	public void reset(){
		acn = new LinkedList<HyperAlertCorrelationNode>();
		PredicateFactToNode = new HashMap<Key,List<HyperAlertCorrelationNode>>();
		hyperAlerts = new LinkedList<HyperAlert>();
	}
	public void run(List<HyperAlert> hyperAlerts){

		ListIterator<HyperAlert> hyperAlertsIter = hyperAlerts.listIterator();
		while(hyperAlertsIter.hasNext()){

			HyperAlert curHyperAlert = hyperAlertsIter.next();
			HyperAlertCorrelationNode node = null;
			List<Predicate> prereqs = curHyperAlert.getType().getPrerequisites();
			if(prereqs == null || prereqs.size() == 0){
				node = addConsequences(curHyperAlert);
				continue;
			}
			
			boolean allPredicatesAreTrue = true;
			ListIterator<Predicate> prerequisiteIter = prereqs.listIterator();
			Set<HyperAlertCorrelationNode> successfulNodes = new HashSet<HyperAlertCorrelationNode>();//temporary set of successful nodes
			
			do{
				Predicate p = prerequisiteIter.next();
				Key key = new Key(curHyperAlert,p);
				
				List<HyperAlertCorrelationNode> nodesWithPredicate = PredicateFactToNode.get(key);
				
				if(nodesWithPredicate == null || nodesWithPredicate.size() == 0){
					allPredicatesAreTrue = false;
					break;
				}else{
					
					ListIterator<HyperAlertCorrelationNode> nodesWithPredicateIter = nodesWithPredicate.listIterator();
					
					int initialMatches = nodesWithPredicate.size();
					while(nodesWithPredicateIter.hasNext()){
						HyperAlertCorrelationNode curNodeWithPredicate = nodesWithPredicateIter.next();

						if(curHyperAlert.getTimestamp().getTime() - curNodeWithPredicate.getAlert().getTimestamp().getTime() < gracePeriodHrs*60*60*1000L){
							successfulNodes.add(curNodeWithPredicate);
						}else{
							nodesWithPredicateIter.remove();
						}
					}
					
					if(initialMatches > nodesWithPredicate.size()){
						PredicateFactToNode.put(key, nodesWithPredicate);
						if(nodesWithPredicate.size() == 0){
							allPredicatesAreTrue = false;
						}
					}
				}
			}while(prerequisiteIter.hasNext() && allPredicatesAreTrue);
			if(allPredicatesAreTrue){
				node = addConsequences(curHyperAlert);

				Iterator<HyperAlertCorrelationNode> successfulNodeIter = successfulNodes.iterator();
				while(successfulNodeIter.hasNext()){
					HyperAlertCorrelationNode curNode = successfulNodeIter.next();
					curNode.addConsequenceNode(node);
					node.addPrerequisiteNode(curNode);
				}
			}
		}
		if(debug){
			System.out.println("Alert Correlation List Size " + acn.size());
			for(int i = 0;i<acn.size();i++){
				System.out.println("Alert " + i +": "+acn.get(i).getAlert().getType().getName());
			}
			System.out.println("Ending Correlation");
		}
		
	}
	
	private HyperAlertCorrelationNode addConsequences(HyperAlert hyperAlert){

		List<Predicate> consequences = hyperAlert.getType().getConsequences();
		
		HyperAlertCorrelationNode node = new HyperAlertCorrelationNode(hyperAlert);
		
		if(consequences == null || consequences.size() == 0){
			if(hyperAlert.getType().getPrerequisites().size() == 0){
				return null;
			}
		}else{
			
			ListIterator<Predicate> consequenceIter = consequences.listIterator();
			while(consequenceIter.hasNext()){
				Predicate consequence = consequenceIter.next();
				for(Implication implication : implications){
					if(consequence.equals(implication.getImplyingElement())){
						for(Predicate p : implication.getImpliedElements()){
							if(!consequences.contains(p)){
								consequenceIter.add(p);
							}
						}
					}
				}
				Key key = new Key(hyperAlert,consequence);
				
				List<HyperAlertCorrelationNode> alertsWithConsequence;
				if(PredicateFactToNode.containsKey(key)){
					alertsWithConsequence = PredicateFactToNode.get(key);
				}else{
					alertsWithConsequence = new LinkedList<HyperAlertCorrelationNode>();
				}
				alertsWithConsequence.add(node);
				PredicateFactToNode.put(key, alertsWithConsequence);
			}
		}
		acn.add(node);
		return node;
	}
	
	public List<HyperAlert> getHyperAlerts() {
		return hyperAlerts;
	}
	
	public List<HyperAlert> get() {
		return hyperAlerts;
	}
	public List<HyperAlertCorrelationNode> getHyperAlertCorrelationNodes(){
		return acn;
	}
	public static void main(String[] args) {
		KnowledgeBase kb = KnowledgeBase.getInstance();
		kb.load();
		List<HyperAlert> alerts = new LinkedList<HyperAlert>();
		alerts.add(new HyperAlert(kb.getHyperAlertType(408),new Date(),111222333444L, 1111, 123123123123L, 1234));
		alerts.add(new HyperAlert(kb.getHyperAlertType(585),new Date(),111222333444L, 1111, 123123123123L, 1234));
		CorrelationEngine ce = new CorrelationEngine();
		ce.run(alerts);
	}
}
